#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
量比计算器

支持两种量比计算方式：
1. 累计量比（固定窗口实时量比）：基于累计成交量和交易分钟数
2. 连续量比（滑动窗口量比）：基于逐笔数据的滑动窗口

设计原则：
- 统一接口，供分时回测、实盘、模拟盘、监控使用
- 支持缓存机制，避免重复计算
- 支持配置化，灵活切换计算方式
"""

import numpy as np
import logging
from typing import Dict, Optional, List, Tuple
from collections import deque
from datetime import datetime, timedelta

logger = logging.getLogger(__name__)


class VolumeRatioCalculator:
    """量比计算器"""
    
    @staticmethod
    def get_required_history_days(config: Dict) -> int:
        """
        获取所需的最小历史数据天数（用于自动延长数据加载起始日期）
        
        Args:
            config: 配置字典，包含volume_ratio_calculator配置
        
        Returns:
            int: 所需的最小历史数据天数（自然日，考虑交易日约1.5倍）
        """
        vr_config = config.get('volume_ratio_calculator', {})
        enabled = vr_config.get('enabled', True)
        
        if not enabled:
            return 60  # 默认历史数据需求
        
        # 获取历史数据周期
        cumulative_config = vr_config.get('cumulative', {})
        lookback_days = cumulative_config.get('lookback_days', 5)
        
        # 考虑交易日与自然日的比例（约1.5倍），加上额外余量（10天）
        required_days = int(lookback_days * 1.5) + 10
        
        return required_days
    
    def __init__(self, config: Dict):
        """
        初始化量比计算器
        
        Args:
            config: 配置字典，包含volume_ratio_calculator配置
        """
        self.config = config.get('volume_ratio_calculator', {})
        self.enabled = self.config.get('enabled', True)
        
        # 默认计算方法
        self.default_method = self.config.get('default_method', 'continuous')  # 'cumulative' or 'continuous'
        self.fallback_method = self.config.get('fallback_method', 'cumulative')  # 数据不可用时回退
        
        # 累计量比配置
        cumulative_config = self.config.get('cumulative', {})
        self.cumulative_lookback_days = cumulative_config.get('lookback_days', 5)
        self.daily_trading_minutes = cumulative_config.get('daily_trading_minutes', 240)
        
        # 连续量比配置
        continuous_config = self.config.get('continuous', {})
        self.window_size_seconds = continuous_config.get('window_size_seconds', 300)  # 默认5分钟
        self.min_window_size = continuous_config.get('min_window_size', 60)  # 最小1分钟
        self.max_window_size = continuous_config.get('max_window_size', 900)  # 最大15分钟
        self.adaptive_window = continuous_config.get('adaptive_window', True)  # 是否根据交易时间自动调整
        
        # 缓存配置
        cache_config = self.config.get('cache', {})
        self.cache_enabled = cache_config.get('enabled', True)
        self.cache_ttl_seconds = cache_config.get('ttl_seconds', 60)
        
        # ⭐ 统一缓存（key: (stock, method), value: (volume_ratio, timestamp)）
        self._ratio_cache: Dict[Tuple[str, str], Tuple[float, datetime]] = {}
        
        # ⭐ 滑动窗口缓存（key: stock, value: deque of (timestamp, volume)）
        self._sliding_windows: Dict[str, deque] = {}
        
        logger.info(f"量比计算器初始化: enabled={self.enabled}, default_method={self.default_method}")
    
    def calculate_trading_minutes(self, time_tag: str) -> int:
        """
        计算当日累计开市时间（分钟）
        A股交易时间：9:30-11:30 (上午2小时), 13:00-15:00 (下午2小时)
        
        Args:
            time_tag: 时间标签，格式如 "2025-01-24 10:12:00"
        
        Returns:
            int: 累计交易分钟数
        """
        try:
            if not time_tag:
                return 0
            
            # 解析时间标签
            time_parts = time_tag.split(' ')
            if len(time_parts) >= 2:
                time_str = time_parts[1]
                hour, minute = map(int, time_str.split(':')[:2])
                
                # 使用统一公式计算
                if hour < 11 or (hour == 11 and minute <= 30):
                    # 上午盘：9:30-11:30
                    trading_minutes = (hour - 9) * 60 + minute - 30
                elif hour >= 13:
                    # 下午盘：13:00-15:00
                    trading_minutes = 120 + (hour - 13) * 60 + minute
                else:
                    # 中午休市时间（11:30-13:00），使用上午的120分钟
                    trading_minutes = 120
                
                return max(1, trading_minutes)
            
            return 0
        except Exception as e:
            logger.debug(f"时间解析失败: {e}, time_tag={time_tag}")
            return 0
    
    def calculate_cumulative_volume_ratio(
        self,
        stock: str,
        current_volume: float,
        historical_daily_data,
        actual_trading_minutes: Optional[int] = None,
        time_tag: Optional[str] = None
    ) -> float:
        """
        计算累计量比（固定窗口实时量比）
        
        量比公式：当前累计成交量 ÷ (过去5日平均每分钟成交量 × 当前已交易分钟数)
        
        Args:
            stock: 股票代码
            current_volume: 当前累计成交量（从开盘到当前）
            historical_daily_data: 历史日线数据（DataFrame，包含volume列）
            actual_trading_minutes: 实际交易分钟数（可选，如果不提供则从time_tag计算）
            time_tag: 时间标签（可选，用于计算交易分钟数）
        
        Returns:
            float: 量比值（>0表示成功，0表示失败）
        """
        try:
            # 计算实际交易分钟数
            if actual_trading_minutes is None:
                if time_tag:
                    actual_trading_minutes = self.calculate_trading_minutes(time_tag)
                else:
                    logger.warning(f"{stock} 无法计算交易分钟数：缺少time_tag或actual_trading_minutes")
                    return 0.0
            
            # 获取过去N日的平均成交量
            if historical_daily_data is not None and len(historical_daily_data) >= self.cumulative_lookback_days:
                recent_volumes = historical_daily_data['volume'].iloc[-self.cumulative_lookback_days:]
                avg_volume_nd = recent_volumes.mean()
                
                # 计算过去N日平均每分钟成交量
                past_nd_avg_per_minute = avg_volume_nd / self.daily_trading_minutes
                
                # 使用实际交易分钟数（从开盘到当前的分钟数）
                current_trading_minutes = min(actual_trading_minutes, self.daily_trading_minutes)
                current_trading_minutes = max(1, current_trading_minutes)  # 至少1分钟
                
                # 计算量比
                expected_volume = past_nd_avg_per_minute * current_trading_minutes
                volume_ratio = current_volume / expected_volume if expected_volume > 0 else 0
            else:
                # 如果历史数据不足，使用前一日成交量
                if historical_daily_data is not None and len(historical_daily_data) > 0:
                    prev_volume = historical_daily_data['volume'].iloc[-1]
                    volume_ratio = current_volume / prev_volume if prev_volume > 0 else 0
                else:
                    logger.debug(f"{stock} 历史数据不足，无法计算量比")
                    return 0.0
            
            return round(volume_ratio, 2)
        except Exception as e:
            logger.warning(f"计算{stock}累计量比失败: {e}")
            return 0.0
    
    def _get_adaptive_window_size(self, current_time: datetime) -> int:
        """
        根据交易时间获取自适应窗口大小
        
        Args:
            current_time: 当前时间
        
        Returns:
            int: 窗口大小（秒）
        """
        if not self.adaptive_window:
            return self.window_size_seconds
        
        hour = current_time.hour
        minute = current_time.minute
        
        # 开盘30分钟内：窗口5分钟（捕捉快速变化）
        if hour == 9 and minute >= 30:
            return 300  # 5分钟
        elif hour == 10 and minute < 30:
            return 300  # 5分钟
        
        # 30-120分钟：窗口10分钟（稳定判断）
        elif (hour == 10 and minute >= 30) or (hour == 11 and minute <= 30):
            return 600  # 10分钟
        
        # 120分钟后：窗口15分钟（避免噪音）
        else:
            return 900  # 15分钟
    
    def update_sliding_window(self, stock: str, timestamp: datetime, volume: float):
        """
        更新滑动窗口（添加新的逐笔数据）
        
        Args:
            stock: 股票代码
            timestamp: 时间戳
            volume: 成交量
        """
        if stock not in self._sliding_windows:
            self._sliding_windows[stock] = deque()
        
        window = self._sliding_windows[stock]
        window_size = self._get_adaptive_window_size(timestamp)
        
        # 添加新数据
        window.append((timestamp, volume))
        
        # 移除过期数据（超过窗口大小）
        cutoff_time = timestamp - timedelta(seconds=window_size)
        while window and window[0][0] < cutoff_time:
            window.popleft()
    
    def calculate_continuous_volume_ratio(
        self,
        stock: str,
        tick_volumes: Optional[List[Tuple[datetime, float]]] = None,
        historical_daily_data=None,
        current_time: Optional[datetime] = None
    ) -> float:
        """
        计算连续量比（滑动窗口量比）
        
        量比公式：最近N秒成交量 ÷ (过去5日平均每秒成交量 × N秒)
        
        Args:
            stock: 股票代码
            tick_volumes: 逐笔成交量列表（可选，如果不提供则使用滑动窗口缓存）
            historical_daily_data: 历史日线数据（DataFrame，包含volume列）
            current_time: 当前时间（可选，用于自适应窗口）
        
        Returns:
            float: 量比值（>0表示成功，0表示失败）
        """
        try:
            # 获取窗口数据
            if tick_volumes is not None:
                # 使用提供的逐笔数据
                if not tick_volumes:
                    return 0.0
                
                # 获取窗口大小
                if current_time is None:
                    current_time = tick_volumes[-1][0] if tick_volumes else datetime.now()
                
                window_size = self._get_adaptive_window_size(current_time)
                window_size = max(self.min_window_size, min(window_size, self.max_window_size))
                
                # 计算窗口内的成交量
                cutoff_time = current_time - timedelta(seconds=window_size)
                recent_volumes = [vol for ts, vol in tick_volumes if ts >= cutoff_time]
                recent_volume = sum(recent_volumes)
            else:
                # 使用滑动窗口缓存
                if stock not in self._sliding_windows:
                    logger.debug(f"{stock} 滑动窗口为空，无法计算连续量比")
                    return 0.0
                
                window = self._sliding_windows[stock]
                if not window:
                    return 0.0
                
                # 获取窗口大小
                if current_time is None:
                    current_time = window[-1][0] if window else datetime.now()
                
                window_size = self._get_adaptive_window_size(current_time)
                window_size = max(self.min_window_size, min(window_size, self.max_window_size))
                
                # 计算窗口内的成交量
                cutoff_time = current_time - timedelta(seconds=window_size)
                recent_volumes = [vol for ts, vol in window if ts >= cutoff_time]
                recent_volume = sum(recent_volumes)
            
            if recent_volume <= 0:
                return 0.0
            
            # 获取历史数据
            if historical_daily_data is None or len(historical_daily_data) < self.cumulative_lookback_days:
                logger.debug(f"{stock} 历史数据不足，无法计算连续量比")
                return 0.0
            
            # 计算过去N日平均每日成交量
            recent_volumes = historical_daily_data['volume'].iloc[-self.cumulative_lookback_days:]
            avg_volume_nd = recent_volumes.mean()
            
            # 计算过去N日平均每秒成交量
            daily_trading_seconds = self.daily_trading_minutes * 60
            past_nd_avg_per_second = avg_volume_nd / daily_trading_seconds
            
            # 计算量比
            expected_volume = past_nd_avg_per_second * window_size
            volume_ratio = recent_volume / expected_volume if expected_volume > 0 else 0
            
            return round(volume_ratio, 2)
        except Exception as e:
            logger.warning(f"计算{stock}连续量比失败: {e}")
            return 0.0
    
    def get_volume_ratio(
        self,
        stock: str,
        method: Optional[str] = None,
        **kwargs
    ) -> float:
        """
        获取量比（统一入口）
        
        Args:
            stock: 股票代码
            method: 计算方法（'cumulative' 或 'continuous'），如果不提供则使用默认方法
            **kwargs: 其他参数，根据计算方法不同而不同
        
        Returns:
            float: 量比值（>0表示成功，0表示失败）
        """
        if not self.enabled:
            return 0.0
        
        method = method or self.default_method
        
        # 检查缓存
        if self.cache_enabled:
            cache_key = (stock, method)
            if cache_key in self._ratio_cache:
                cached_ratio, cached_time = self._ratio_cache[cache_key]
                if (datetime.now() - cached_time).total_seconds() < self.cache_ttl_seconds:
                    return cached_ratio
        
        try:
            # 根据方法计算量比
            if method == 'continuous':
                volume_ratio = self.calculate_continuous_volume_ratio(stock, **kwargs)
            else:
                volume_ratio = self.calculate_cumulative_volume_ratio(stock, **kwargs)
            
            # 如果连续量比失败，尝试回退到累计量比
            if volume_ratio <= 0 and method == 'continuous' and self.fallback_method == 'cumulative':
                logger.debug(f"{stock} 连续量比计算失败，回退到累计量比")
                volume_ratio = self.calculate_cumulative_volume_ratio(stock, **kwargs)
            
            # 更新缓存
            if self.cache_enabled and volume_ratio > 0:
                cache_key = (stock, method)
                self._ratio_cache[cache_key] = (volume_ratio, datetime.now())
            
            return volume_ratio
        except Exception as e:
            logger.warning(f"{stock} 获取量比失败: {e}")
            # 尝试回退方法
            if method != self.fallback_method:
                try:
                    return self.get_volume_ratio(stock, method=self.fallback_method, **kwargs)
                except:
                    pass
            return 0.0
    
    def clear_cache(self, stock: Optional[str] = None):
        """
        清除缓存
        
        Args:
            stock: 股票代码（如果提供则只清除该股票的缓存，否则清除所有缓存）
        """
        if stock:
            keys_to_remove = [key for key in self._ratio_cache.keys() if key[0] == stock]
            for key in keys_to_remove:
                del self._ratio_cache[key]
            if stock in self._sliding_windows:
                del self._sliding_windows[stock]
        else:
            self._ratio_cache.clear()
            self._sliding_windows.clear()

